//
// Created by ubuntu on 2024/1/3.
//

#include "split_and_merge.h"
#include <algorithm>
#include <Eigen/Dense>
#include <iostream>

namespace line_extraction
{
SplitAndMerge::SplitAndMerge() {}

SplitAndMerge::~SplitAndMerge() {}

void SplitAndMerge::extractLines(std::vector<Line> &lines)
{
    // Resets
    filtered_indices_ = cs_data_.index;
    lines_.clear();
    // 过滤器索引
    filterCloseAndFarPoints();
    // 过滤离群点
    filterOutlierPoints();
    // 如果没有留下足够的点，则不返回行
    if (filtered_indices_.size() <= std::max(params_.min_line_points, 3)) return;
    // 将索引拆分成行，并过滤出短行和稀疏行
    split(filtered_indices_);
    filterLines();
    // 使用最小二乘法拟合每条直线并合并共线
    for (auto &line : lines_) line.leastSqFit();
    // 如果有多行，请检查是否应根据合并条件合并行
    if (lines_.size() > 1) mergeLines();
    lines = lines_;
}

///////////////////////////////////////////////////////////////////////////////
// Filtering points
///////////////////////////////////////////////////////////////////////////////
void SplitAndMerge::filterCloseAndFarPoints()
{
    std::vector<unsigned int> output;
    for (unsigned int filtered_indice : filtered_indices_)
    {
        const double &range = range_data_.ranges[filtered_indice];
        if (range >= params_.min_range && range <= params_.max_range) output.push_back(filtered_indice);
    }
    filtered_indices_ = output;
}

void SplitAndMerge::filterOutlierPoints()
{
    if (filtered_indices_.size() < 3) return;
    std::vector<unsigned int> output;
    unsigned int p_i, p_j, p_k;
    for (std::size_t i = 0; i < filtered_indices_.size(); ++i)
    {
        // Get two closest neighbours
        p_i = filtered_indices_[i];
        // first point
        if (i == 0)
        {
            p_j = filtered_indices_[i + 1];
            p_k = filtered_indices_[i + 2];
        }
            // last point
        else if (i == filtered_indices_.size() - 1)
        {
            p_j = filtered_indices_[i - 1];
            p_k = filtered_indices_[i - 2];
        }
            // middle points
        else
        {
            p_j = filtered_indices_[i - 1];
            p_k = filtered_indices_[i + 1];
        }
        // 检查点是否为异常值
        if (fabs(range_data_.ranges[p_i] - range_data_.ranges[p_j]) > params_.outlier_dist &&
            fabs(range_data_.ranges[p_i] - range_data_.ranges[p_k]) > params_.outlier_dist)
        {
            // 检查是否靠近连接其邻居的线路
            std::vector<unsigned int> line_indices;
            line_indices.push_back(p_j);
            line_indices.push_back(p_k);
            Line line(cs_data_, range_data_, params_, line_indices);
            line.endpointFit();
            if (line.distToPoint(p_i) > params_.min_split_dist) continue; // point is an outlier
        }
        output.push_back(p_i);
    }
    filtered_indices_ = output;
}

///////////////////////////////////////////////////////////////////////////////
// Filtering and merging lines
///////////////////////////////////////////////////////////////////////////////
void SplitAndMerge::filterLines()
{
    std::vector<Line> output;
    for (const auto &line : lines_)
    {
        if (line.length() >= params_.min_line_length && line.numPoints() >= params_.min_line_points)
            output.push_back(line);
    }
    lines_ = output;
}

void SplitAndMerge::mergeLines()
{
    std::vector<Line> merged_lines;
    for (std::size_t i = 1; i < lines_.size(); ++i)
    {
        // Get L, P_1, P_2 of consecutive lines
        Eigen::Vector2d L_1(lines_[i - 1].getRadius(), lines_[i - 1].getAngle());
        Eigen::Vector2d L_2(lines_[i].getRadius(), lines_[i].getAngle());
        Eigen::Matrix2d P_1;
        P_1 << lines_[i - 1].getCovariance()[0], lines_[i - 1].getCovariance()[1],
                lines_[i - 1].getCovariance()[2], lines_[i - 1].getCovariance()[3];
        Eigen::Matrix2d P_2;
        P_2 << lines_[i].getCovariance()[0], lines_[i].getCovariance()[1],
                lines_[i].getCovariance()[2], lines_[i].getCovariance()[3];
        // Merge lines if chi-squared distance is less than 3
        if (chiSquared(L_1 - L_2, P_1, P_2) < 3)
        {
            // Get merged angle, radius, and covariance
            Eigen::Matrix2d P_m = (P_1.inverse() + P_2.inverse()).inverse();
            Eigen::Vector2d L_m = P_m * (P_1.inverse() * L_1 + P_2.inverse() * L_2);
            // Populate new line with these merged parameters
            boost::array<double, 4> cov{};
            cov[0] = P_m(0, 0);
            cov[1] = P_m(0, 1);
            cov[2] = P_m(1, 0);
            cov[3] = P_m(1, 1);
            std::vector<unsigned int> index;
            const std::vector<unsigned int> &ind_1 = lines_[i - 1].getIndices();
            const std::vector<unsigned int> &ind_2 = lines_[i].getIndices();
            index.resize(ind_1.size() + ind_2.size());
            index.insert(index.end(), ind_1.begin(), ind_1.end());
            index.insert(index.end(), ind_2.begin(), ind_2.end());
            Line merged_line(L_m[1], L_m[0], cov, lines_[i - 1].getStart(), lines_[i].getEnd(), index);
            // Project the new endpoints
            merged_line.projectEndpoints();
            lines_[i] = merged_line;
        } else merged_lines.push_back(lines_[i - 1]);
        if (i == lines_.size() - 1) merged_lines.push_back(lines_[i]);
    }
    lines_ = merged_lines;
}

///////////////////////////////////////////////////////////////////////////////
// Splitting points into lines
///////////////////////////////////////////////////////////////////////////////
void SplitAndMerge::split(const std::vector<unsigned int> &index)
{
    // 仅在单个点时不拆分（仅在被间隙孤立时发生）
    if (index.size() <= 1) return;
    Line line(cs_data_, range_data_, params_, index);
    line.endpointFit();
    double dist_max = 0;
    double gap_max = 0;
    double dist, gap;
    int i_max, i_gap;
    // 找到最远的点和最大的间隙
    for (std::size_t i = 1; i < index.size() - 1; ++i)
    {
        dist = line.distToPoint(index[i]);
        if (dist > dist_max)
        {
            dist_max = dist;
            i_max = i;
        }
        gap = distBetweenPoints(index[i], index[i + 1]);
        if (gap > gap_max)
        {
            gap_max = gap;
            i_gap = i;
        }
    }
    // Check for gaps at endpoints
    double gap_start = distBetweenPoints(index[0], index[1]);
    if (gap_start > gap_max)
    {
        gap_max = gap_start;
        i_gap = 1;
    }
    double gap_end = distBetweenPoints(index.rbegin()[1], index.rbegin()[0]);
    if (gap_end > gap_max)
    {
        gap_max = gap_end;
        i_gap = index.size() - 1;
    }
    // 检查线段是否符合要求或应拆分
    if (dist_max < params_.min_split_dist && gap_max < params_.max_line_gap) lines_.push_back(line);
    else
    {
        int i_split = dist_max >= params_.min_split_dist ? i_max : i_gap;
        std::vector<unsigned int> first_split(&index[0], &index[i_split - 1]);
        std::vector<unsigned int> second_split(&index[i_split], &index.back());
        split(first_split);
        split(second_split);
    }
}

///////////////////////////////////////////////////////////////////////////////
// Utility methods
///////////////////////////////////////////////////////////////////////////////
double SplitAndMerge::chiSquared(const Eigen::Vector2d &dL, const Eigen::Matrix2d &P_1,
                                  const Eigen::Matrix2d &P_2)
{
    return dL.transpose() * (P_1 + P_2).inverse() * dL;
}

double SplitAndMerge::distBetweenPoints(unsigned int index_1, unsigned int index_2)
{
    return sqrt(pow(range_data_.xs[index_1] - range_data_.xs[index_2], 2) +
                pow(range_data_.ys[index_1] - range_data_.ys[index_2], 2));
}

} // namespace line_extraction

